import sys
import argparse
import os

 
from ultralytics import YOLO



if __name__ == '__main__':
    '''
    # 加载模型
    model = YOLO("yolov8n.yaml")  # 从头开始构建新模型
    model = YOLO("yolov8n.pt")  # 加载预训练模型（建议用于训练）

    # 使用模型
    model.train(data="red_2.yaml", epochs = num_epochs)  # 训练模型
    metrics = model.val()  # 在验证集上评估模型性能
    results = model("bus.jpg")  # 对图像进行预测
    #success = model.export(format="onnx")  # 将模型导出为 ONNX 格式
    '''
    parser = argparse.ArgumentParser()
    parser.add_argument('--yolo_weights', type=str, default='yolo11n.pt', help='model.pt path')
    parser.add_argument('--dataset_yaml', type=str, default='coco128.yaml', help='dataset yaml')
    parser.add_argument('--num_epochs', type=int, default=10, help='num epochs')
    args = parser.parse_args()
     
    # 加载模型
    #model = YOLO("yolov8n.yaml")  # 从头开始构建新模型
    model = YOLO(args.yolo_weights)  # 加载预训练模型（建议用于训练）

    # 使用模型
    model.train(data=args.dataset_yaml, epochs = args.num_epochs)  # 训练模型
    #metrics = model.val()  # 在验证集上评估模型性能
    #results = model("bus.jpg")  # 对图像进行预测
    #success = model.export(format="onnx")  # 将模型导出为 ONNX 格式
    